package server

import (
	"context"
	"fmt"
	"log"
	"math/big"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"

	"github.com/google/uuid"
	"github.com/joho/godotenv"
	grpc_err "github.com/linchao/zkp_auth/api/v2/err"
	api "github.com/linchao/zkp_auth/api/v2/proto"
	cp_zkp "github.com/linchao/zkp_auth/internal/cpzkp"
	"github.com/linchao/zkp_auth/lib/util"
	"github.com/linchao/zkp_auth/models"
	"google.golang.org/grpc"
)

type CPZKP interface {
	InitCPZKPParams() (*cp_zkp.CPZKPParams, error)
}

type Config struct {
	CPZKP CPZKP
	// 新增DB支持
	DBPath string
}

type RegParams struct {
	y1 *big.Int
	y2 *big.Int
}

type AuthParams struct {
	user string
	c    *big.Int
	r1   *big.Int
	r2   *big.Int
}

type grpcServer struct {
	api.UnimplementedAuthServer

	// 依然保留内存方式，但主要用于认证过程中的临时存储
	// 设备注册信息已转移到SQLite数据库
	RegDir map[string]RegParams

	// 认证过程中的临时信息
	AuthDir map[string]AuthParams

	*Config
}

func RunServer(config *Config) {
	err := godotenv.Load(".env")
	if err != nil {
		log.Printf("warning: error loading .env file: %v", err)
	}

	// 初始化数据库
	dbPath := config.DBPath
	if dbPath == "" {
		// 如果未指定路径，使用默认路径
		dbPath = "./data/auth.db"
	}

	// 确保数据目录存在
	dir := filepath.Dir(dbPath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		log.Fatalf("failed to create data directory: %v", err)
	}

	// 初始化数据库
	if err := models.InitDB(dbPath); err != nil {
		log.Fatalf("failed to initialize database: %v", err)
	}
	defer models.CloseDB()

	grpcServerAddr := os.Getenv("SERVER_ADDRESS")
	if grpcServerAddr == "" {
		grpcServerAddr = ":50051" // 默认端口
	}

	listener, err := net.Listen("tcp", grpcServerAddr)
	if err != nil {
		log.Fatalf("failed to dial server: %v", err)
	}

	// Create a new gRPC server and register the service
	grpcServer, err := NewGRPCSever(config)
	if err != nil {
		log.Fatalf("failed to create gRPC server: %v", err)
	}

	// 信号处理
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

	done := make(chan struct{})

	go func() {
		<-quit
		log.Println("收到终止信号，正在优雅关闭gRPC服务...")
		grpcServer.GracefulStop()
		cwd, _ := os.Getwd()
		log.Printf("当前工作目录: %s", cwd)
		files, err := filepath.Glob("./token/*.json")
		if err != nil {
			log.Printf("查找token目录下json文件失败: %v", err)
		} else {
			for _, f := range files {
				if err := os.Remove(f); err != nil {
					log.Printf("删除文件%s失败: %v", f, err)
				} else {
					log.Printf("已删除文件: %s", f)
				}
			}
		}
		close(done)
	}()

	log.Printf("grpc server listening on: %s\n", listener.Addr().String())
	log.Printf("database initialized at: %s\n", dbPath)

	// Start the gRPC server
	if err := grpcServer.Serve(listener); err != nil {
		log.Fatalf("failed to start gRPC server: %v", err)
	}

	// 等待信号处理完成
	<-done
	log.Println("Server shutdown complete.")
}

func newgrpcServer(config *Config) (*grpcServer, error) {
	// initialize the server with ZKP system params and an empty user directory
	return &grpcServer{
		RegDir:  make(map[string]RegParams),
		AuthDir: make(map[string]AuthParams),
		Config:  config,
	}, nil
}

// NewGRPCSever: creates a grpc server and registers the service to that server
func NewGRPCSever(config *Config) (*grpc.Server, error) {
	gsrv := grpc.NewServer()
	srv, err := newgrpcServer(config)
	if err != nil {
		return nil, err
	}
	api.RegisterAuthServer(gsrv, srv)
	return gsrv, nil
}

// Register: 注册新的设备，将y1和y2保存到数据库
func (s *grpcServer) Register(ctx context.Context, req *api.RegisterRequest) (
	*api.RegisterResponse, error) {

	// 转换成deviceID，更适合IoT场景
	deviceID := req.User
	if deviceID == "" {
		return nil, fmt.Errorf("device ID不能为空")
	}

	// 首先检查设备是否已经注册
	device, err := models.GetDevice(ctx, deviceID)
	if err != nil {
		return nil, fmt.Errorf("database error: %v", err)
	}

	if device != nil {
		return nil, grpc_err.ErrInvalidRegistration{User: deviceID}
	}

	// 验证y1和y2参数
	if req.Y1 == "" || req.Y2 == "" {
		return nil, fmt.Errorf("Y1和Y2参数不能为空")
	}

	// 解析大整数
	Y1, err := util.ParseBigInt(req.Y1, "y1")
	if err != nil {
		return nil, err
	}

	Y2, err := util.ParseBigInt(req.Y2, "y2")
	if err != nil {
		return nil, err
	}

	// 保存到内存（仅用于当前会话）
	s.RegDir[deviceID] = RegParams{
		y1: Y1,
		y2: Y2,
	}

	// 保存到数据库（持久化）
	newDevice := &models.Device{
		DeviceID: deviceID,
		Y1:       req.Y1,
		Y2:       req.Y2,
	}

	if err := models.SaveDevice(ctx, newDevice); err != nil {
		return nil, fmt.Errorf("failed to save device: %v", err)
	}

	log.Printf("Device %s registered successfully", deviceID)
	return &api.RegisterResponse{}, nil
}

func (s *grpcServer) CreateAuthenticationChallenge(ctx context.Context, req *api.AuthenticationChallengeRequest) (
	*api.AuthenticationChallengeResponse, error) {

	deviceID := req.User
	if deviceID == "" {
		return nil, fmt.Errorf("device ID不能为空")
	}

	// 验证r1和r2参数
	if req.R1 == "" || req.R2 == "" {
		return nil, fmt.Errorf("R1和R2参数不能为空")
	}

	// 从数据库检查设备是否已注册
	device, err := models.GetDevice(ctx, deviceID)
	if err != nil {
		return nil, fmt.Errorf("database error: %v", err)
	}

	if device == nil {
		return nil, fmt.Errorf("device %s is not registered on the server", deviceID)
	}

	// 将设备信息加载到内存中（如果不存在）
	if _, exists := s.RegDir[deviceID]; !exists {
		Y1, err := util.ParseBigInt(device.Y1, "y1")
		if err != nil {
			return nil, err
		}

		Y2, err := util.ParseBigInt(device.Y2, "y2")
		if err != nil {
			return nil, err
		}

		s.RegDir[deviceID] = RegParams{
			y1: Y1,
			y2: Y2,
		}
	}

	cpzkpParams, err := s.Config.CPZKP.InitCPZKPParams()
	if err != nil {
		return nil, err
	}

	// 创建验证者并生成挑战
	verifier := &cp_zkp.Verifier{}
	c, err := verifier.CreateProofChallenge(cpzkpParams)
	if err != nil {
		return nil, err
	}

	// 生成认证ID
	authID, err := uuid.NewRandom()
	if err != nil {
		return nil, err
	}

	R1, err := util.ParseBigInt(req.R1, "r1")
	if err != nil {
		return nil, err
	}

	R2, err := util.ParseBigInt(req.R2, "r2")
	if err != nil {
		return nil, err
	}

	auth_id := authID.String()
	s.AuthDir[auth_id] = AuthParams{
		user: deviceID,
		c:    c,
		r1:   R1,
		r2:   R2,
	}

	return &api.AuthenticationChallengeResponse{
		AuthId: auth_id,
		C:      c.String(),
	}, nil
}

func (s *grpcServer) VerifyAuthentication(ctx context.Context, req *api.AuthenticationAnswerRequest) (
	*api.AuthenticationAnswerResponse, error) {

	if req.AuthId == "" {
		return nil, fmt.Errorf("authentication ID不能为空")
	}

	if req.S == "" {
		return nil, fmt.Errorf("response parameter S不能为空")
	}

	// 检查认证ID是否有效
	authParams, idExists := s.AuthDir[req.AuthId]
	if !idExists {
		return nil, fmt.Errorf("invalid authentication id: %s specified", req.AuthId)
	}

	cpzkpParams, err := s.Config.CPZKP.InitCPZKPParams()
	if err != nil {
		return nil, err
	}

	// 获取用户名和挑战值
	deviceID := authParams.user
	c := authParams.c

	// 获取认证参数
	regParams, exists := s.RegDir[deviceID]
	if !exists {
		// 如果内存中没有，尝试从数据库加载
		device, err := models.GetDevice(ctx, deviceID)
		if err != nil {
			return nil, fmt.Errorf("database error: %v", err)
		}
		if device == nil {
			return nil, fmt.Errorf("device %s not found", deviceID)
		}

		y1, err := util.ParseBigInt(device.Y1, "y1")
		if err != nil {
			return nil, err
		}

		y2, err := util.ParseBigInt(device.Y2, "y2")
		if err != nil {
			return nil, err
		}

		// 保存到内存
		s.RegDir[deviceID] = RegParams{
			y1: y1,
			y2: y2,
		}
		regParams = s.RegDir[deviceID]
	}

	y1 := regParams.y1
	y2 := regParams.y2
	r1 := authParams.r1
	r2 := authParams.r2

	// 解析响应值
	S, err := util.ParseBigInt(req.S, "s")
	if err != nil {
		return nil, err
	}

	// 验证证明
	verifier := &cp_zkp.Verifier{}
	isValidProof := verifier.VerifyProof(y1, y2, r1, r2, c, S, cpzkpParams)
	if !isValidProof {
		return nil, grpc_err.ErrInvalidChallengeResponse{S: req.S}
	}

	// 生成会话ID
	sessionID, err := uuid.NewRandom()
	if err != nil {
		return nil, err
	}

	// 将会话ID保存到数据库
	sessionIDStr := sessionID.String()
	device, err := models.GetDevice(ctx, deviceID)
	if err != nil {
		log.Printf("获取设备信息失败: %v", err)
		// 这里即使出错也不影响认证流程，用户仍然能够通过认证
	} else if device != nil {
		// 更新设备的会话ID
		device.SessionID = sessionIDStr
		if err := models.SaveDevice(ctx, device); err != nil {
			log.Printf("保存会话ID失败: %v", err)
			// 这里即使出错也不影响认证流程，用户仍然能够通过认证
		} else {
			log.Printf("已将会话ID %s 保存到设备 %s 的记录中", sessionIDStr, deviceID)
		}
	}

	log.Printf("Device %s authenticated successfully", deviceID)
	return &api.AuthenticationAnswerResponse{SessionId: sessionIDStr}, nil
}
